from hera.workflows import DAG, Container, Task, Workflow
from hera.workflows.models import Arguments, Inputs, Item, Outputs, Parameter, ValueFrom

with Workflow(
    api_version="argoproj.io/v1alpha1",
    kind="Workflow",
    generate_name="parameter-aggregation-dag-",
    entrypoint="parameter-aggregation",
) as w:
    with DAG(
        name="parameter-aggregation",
    ) as invocator:
        Task(
            with_items=[
                Item(
                    __root__=1,
                ),
                Item(
                    __root__=2,
                ),
                Item(
                    __root__=3,
                ),
                Item(
                    __root__=4,
                ),
            ],
            arguments=Arguments(
                parameters=[
                    Parameter(
                        name="num",
                        value="{{item}}",
                    )
                ],
            ),
            name="odd-or-even",
            template="odd-or-even",
        )
        Task(
            arguments=Arguments(
                parameters=[
                    Parameter(
                        name="message",
                        value="{{tasks.odd-or-even.outputs.parameters.num}}",
                    )
                ],
            ),
            name="print-nums",
            template="print-message",
            depends="odd-or-even",
        )
        Task(
            arguments=Arguments(
                parameters=[
                    Parameter(
                        name="message",
                        value="{{tasks.odd-or-even.outputs.parameters.evenness}}",
                    )
                ],
            ),
            name="print-evenness",
            template="print-message",
            depends="odd-or-even",
        )
        Task(
            with_param="{{tasks.odd-or-even.outputs.parameters}}",
            arguments=Arguments(
                parameters=[
                    Parameter(
                        name="num",
                        value="{{item.num}}",
                    )
                ],
            ),
            name="divide-by-2",
            template="divide-by-2",
            when="{{item.evenness}} == even",
            depends="odd-or-even",
        )
        Task(
            with_param="{{tasks.divide-by-2.outputs.result}}",
            arguments=Arguments(
                parameters=[
                    Parameter(
                        name="message",
                        value="{{item}}",
                    )
                ],
            ),
            name="print",
            template="print-message",
            depends="divide-by-2",
        )
    Container(
        inputs=Inputs(
            parameters=[
                Parameter(
                    name="num",
                )
            ],
        ),
        name="odd-or-even",
        outputs=Outputs(
            parameters=[
                Parameter(
                    name="num",
                    value_from=ValueFrom(
                        path="/tmp/num",
                    ),
                ),
                Parameter(
                    name="evenness",
                    value_from=ValueFrom(
                        path="/tmp/even",
                    ),
                ),
            ],
        ),
        args=[
            'sleep 1 &&\necho {{inputs.parameters.num}} > /tmp/num &&\nif [ $(({{inputs.parameters.num}}%2)) -eq 0 ]; then\n  echo "even" > /tmp/even;\nelse\n  echo "odd" > /tmp/even;\nfi\n'
        ],
        command=["sh", "-c"],
        image="alpine:latest",
    )
    Container(
        inputs=Inputs(
            parameters=[
                Parameter(
                    name="num",
                )
            ],
        ),
        name="divide-by-2",
        args=["echo $(({{inputs.parameters.num}}/2))"],
        command=["sh", "-c"],
        image="alpine:latest",
    )
    Container(
        inputs=Inputs(
            parameters=[
                Parameter(
                    name="message",
                )
            ],
        ),
        name="print-message",
        args=["{{inputs.parameters.message}}"],
        command=["echo"],
        image="busybox",
    )
